#!/usr/bin/env python

from amqplib import client_0_8 as amqp
import yaml

class Producer(object):
    def __init__(self, exchange_name, host, userid='guest', password='guest', virtual_host="/", insist=False):
        """
        Constructor. Initiate connection with the RabbitMQ server.
      
        @param exchange_name name of the exchange to send messages to
        @param host RabbitMQ server host
        @param userid RabbitMQ server username
        @param password RabbitMQ server user's password
        @param virtual_host name of the virtual server to use,defaults to "/"
        @param insist insist on connecting to the specified server, defaults to False
        """
        self.exchange_name = exchange_name
        self.connection = amqp.Connection(host=host,
                                          userid=userid,
                                          password=password,
                                          virtual_host=virtual_host,
                                          insist=insist)
        self.channel = self.connection.channel()

    def __serialize(self, message):
        '''Look at message contents and serialize as desired, then return
        an amqp.Message with the body and content_type set accordingly'''
        
        if type(message).__name__ in ('dict','list','tuple','instance'):
            msg = amqp.Message(yaml.dump(message))
            msg.properties["content_type"] = "text/yaml"
        else:
            msg = amqp.Message(message)
            msg.properties["content_type"] = "text/plain"
        
        return msg
    
    def publish(self, message, routing_key):
        """
        Publish message to exchange using routing key
      
        @param text message to publish
        @param routing_key message routing key
        """
        
        msg = self.__serialize(message)
        msg.properties["delivery_mode"] = 2
        self.channel.basic_publish(exchange=self.exchange_name,
                         routing_key=routing_key, msg=msg)

    def close(self):
        """
        Close channel and connection
        """
        self.channel.close()
        self.connection.close()


class Consumer(object):
    def __init__(self, host, userid='guest', password='guest', virtual_host="/", insist=False):
        """
        Constructor. Initiate a connection to the RabbitMQ server.
      
        @param host RabbitMQ server host
        @param userid RabbitMQ server username
        @param password RabbitMQ server user's password
        @param virtual_host name of the virtual server to use,defaults to "/"
        @param insist insist on connecting to the specified server, defaults to False
        """
        self.connection = amqp.Connection(host=host, userid=userid, password=password, virtual_host="/", insist=False)
        self.channel = self.connection.channel()
        self.callbacks = {}


    def close(self):
        """
        Close channel and connection
        """
        self.channel.close()
        self.connection.close()


    def declare_exchange(self, exchange_name, durable=True, auto_delete=False):
        """
        Create exchange.
      
        @param exchange_name name of the exchange
        @param durable will the server survive a server restart
        @param auto_delete should the server delete the exchange when it is
        no longer in use
        """
        self.exchange_name = exchange_name
        self.channel.exchange_declare(exchange=self.exchange_name, type='topic',
                                      durable=durable, auto_delete=auto_delete)


    def declare_queue(self, queue_name, routing_key, durable=True, exclusive=False, auto_delete=False):
        """
        Create a queue and bind it to the exchange.
      
        @param queue_name Name of the queue to create
        @param routing_key binding key
        @param durable will the queue service a server restart
        @param exclusive only 1 client can work with it
        @param auto_delete should the server delete the exchange when it is
         no longer in use
        """
        self.queue_name = queue_name
        self.routing_key = routing_key
        #channel.queue_bind Returns a tuple containing 3 items:
        #    the name of the queue (essential for automatically-named queues)
        #    message count
        #    consumer count
        (queue_name, msgCount, cnsmrCount) = self.channel.queue_declare(
            queue=self.queue_name, durable=durable,
            exclusive=exclusive, auto_delete=auto_delete)
        self.channel.queue_bind(queue=self.queue_name, exchange=self.exchange_name,
                                routing_key=self.routing_key)
        return (queue_name, msgCount, cnsmrCount)

    def basic_qos(self, prefetch_count=0):
        """
        Create a Quality of Service setting for the channel by setting a
        prefetch count.  If prefetch_count > 0, then fair dispatching is used.
        If prefetch_count=0, then every nth consumer, receives every nth job
        
        @param prefetch_count If prefetch_count == 1, then make sure everyone is
        working all the time, if 0, then pass the workload evenly never mind who
        is starving and who is stuffed.
        
        """
        prefetch_size=0
        a_global=False
        self.channel.basic_qos(prefetch_size=prefetch_size,
                               prefetch_count=prefetch_count,
                               a_global=a_global)

    def start_consuming(self, callback, queue_name=None, consumer_tag='consumer'):
        """
        Start a consumer and register a function to be called when a message is consumed
      
        @param callback function to call
        @param queue_name name of the queue
        @param consumer_tag a client-generated consumer tag to establish context
        """
        if hasattr(self, 'queue_name') or queue_name:
            self.callbacks[consumer_tag] = callback    
            self.channel.basic_consume(queue=getattr(self, 'queue_name', queue_name),
                                     callback=self.__deserialize,
                                     consumer_tag=consumer_tag)
    
    def __deserialize(self, msg):
        '''Do any necessary deserialization necesary, then call the  user
        requested call back'''
        
        if msg.properties['content_type'] == u'text/yaml':
            msg.body = yaml.load(msg.body)
            
        # Now hand the msg object off to the user specified call back
        return self.callbacks[msg.delivery_info['consumer_tag']](msg)
        
    def handle_messages(self):
        """
        Wait for activity on the channel.
        """
        while True:
          self.channel.wait()
